
"""Defines the neural network, loss function and metrics"""
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml
from torchvision import models

class Net(nn.Module):
    def __init__(self, config):
        super(Net, self).__init__()
        # Load a pre-trained ResNet-34 model
        self.resnet = models.resnet34(pretrained=True)

        # Freeze the ResNet layers
        for param in self.resnet.parameters():
            param.requires_grad = False

        # Replace the last fully connected layer of ResNet-34
        num_ftrs = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(num_ftrs, config['model']['num_classes'])

        self.dropout_rate = config['model']['dropout_rate']

    def forward(self, s):
        """
        This function defines how we use the components of our network to operate on an input batch.

        Args:
            s: (Variable) contains a batch of images, of dimension batch_size x 3 x 64 x 64 .

        Returns:
            out: (Variable) dimension batch_size x num_classes with the log probabilities for the labels of each image.

        Note: the dimensions after each step are provided
        """
        # Pass the input through the ResNet-34 backbone
        s = self.resnet(s)  # batch_size x num_classes

        # Apply dropout
        s = F.dropout(s, p=self.dropout_rate, training=self.training)

        # Apply log softmax on each image's output (this is recommended over applying softmax
        # since it is numerically more stable)
        return F.log_softmax(s, dim=1)



if __name__ == '__main__':
    yaml_path = '../experiments/resnet/config.yaml'
    with open(yaml_path, "r", encoding='utf-8') as f:
        config = yaml.safe_load(f)
    assert os.path.isfile(
        yaml_path), "No yaml configuration file found at {}".format(yaml_path)

    s = torch.randn(32, 3, 64, 64)  # [B, C, 64, 64]

    model = Net(config)

    outputs = model(s)

    print(f"outputs {outputs.dtype} {outputs.shape}")